#include "DataCollect.h"
#include <X11/Xlib.h>
#include <X11/Xutil.h>
/**
 * 待优化：
 * 线程池的优化，1.分配的时候采用队列(解决)
 *                             2.未分配的线程休眠，直到被分配以后唤醒（解决）
 * 内存泄漏问题(解决)
 *
 * 待验证：
 *              采集的图片转为yuv后是否正确？(验证成功)
 *
 * TODO：
 *              分离测试代码，写出最终运行代码(解决)
 *
 *
 *
 */

DataCollect::DataCollect()
{
}

DataCollect::DataCollect(int worderNum, int w, int h, int FrameRate) : worderNum(worderNum),
                                                                       w(w),
                                                                       h(h),
                                                                       FrameRate(FrameRate),
                                                                       status(false)
{
    services.resize(worderNum);
    pool = new HzThreadPool(worderNum);
}

DataCollect::~DataCollect()
{
    delete pool;
}

int DataCollect::init()
{
    cout << "init datacollect...." << endl;
    // 激活线程池
    pool->Invoke();
    // 创建cuda服务(分配cuda显存)
    for (int i = 0; i < worderNum; i++)
    {
        services[i] = new CudaRGB32ToYUV444Service(w, h);
    }
    // 打开设备
    XInitThreads();
    display = XOpenDisplay(NULL);
    if (!display)
    {
        std::cerr << "Error: Unable to open X display." << std::endl;
        return -1;
    }
    cout << "init datacollect successfully!" << endl;
    return 0;
}

int DataCollect::release()
{
    cout << "release datacollect..." << endl;
    pool->Stop();
    XCloseDisplay(display);
    for (int i = 0; i < worderNum; i++)
    {
        delete services[i];
    }
    cout << "release datacollect sucessfully!" << endl;
    return 0;
}

void DataCollect::run()
{
    if (status == true)
        return;
    status = true; // 设置状态
    init();        // 初始化
    int interval = 1000 / FrameRate;
    int Pcid = 0;
    while (status)
    {
        // 开始采集
         while (pool->DispatchWorker(new ThreadWorker(&DataCollect::beginCollect, this, new int(Pcid))) == NULL)
        {
            continue;
        }
        // 控制采集速度
        std::this_thread::sleep_for(std::chrono::milliseconds(interval));
        Pcid++;
    }
    yuvstreams.Clear();
    release();
}

void DataCollect::stop()
{
    if (status == false)
        return;
    status = false;
}

int DataCollect::beginCollect(void *arg)
{
    int pcID = *(int *)arg;
    YuvBuffer *buffer = new YuvBuffer(w * h,(int)( 0.25*w * h),(int)( 0.25*w * h));
    //  1. 获取此时的屏幕截图
    uchar *RGB = getScreen(buffer->collectTime);
    // 2. 利用cuda进行数据转换
    services[pcID % worderNum]->RGB32ToYUV420Interface(RGB, buffer->y, buffer->u, buffer->v, w, h);
    // 3 .将数据存到缓存区
    yuvstreams.Push(buffer);
//    yuvbuffersMutex.lock();
//    yuvbuffers.push(buffer);
//    yuvbuffersMutex.unlock();
    delete (int *)arg; // 传进来的参数是通过new出来的，需要在函数内部释放，防止内存泄漏
    delete RGB;      // getScreen返回的参数是new出来的，需要在函数内部释放，防止内存泄漏
    return 0;
}

uchar *DataCollect::getScreen(long &time)
{
    uchar *RGB = new uchar[w * h * 3];
    auto currentTime = std::chrono::system_clock::now();
    time = std::chrono::time_point_cast<std::chrono::milliseconds>(currentTime).time_since_epoch().count();
    // cout << "采集屏幕截图..."<< "时间戳:" << time << endl;

    int screen_num = DefaultScreen(display);
    int screen_width = DisplayWidth(display, screen_num);   // 宽度
    int screen_height = DisplayHeight(display, screen_num); // 高度

    // 2.获取RGB32数据
    XImage *x_image = XGetImage(display, RootWindow(display, screen_num),
                                0, 0, screen_width, screen_height, AllPlanes, ZPixmap);
    if (!x_image)
    {
        std::cerr << "Error: Unable to capture screen image." << std::endl;
        return nullptr;
    }
    unsigned long red_mask = x_image->red_mask; // 这里的掩码也是4bytes，因为在XImage中ZPixmap中一个像素点也是4bytes，这里的掩码是用来萃取中对应的颜色
    unsigned long green_mask = x_image->green_mask;
    unsigned long blue_mask = x_image->blue_mask;
    int t = 0;
    for (int y = 0; y < h; ++y)
    {
        for (int x = 0; x < w; ++x)
        {
            unsigned long pixel = XGetPixel(x_image, x, y);
            uint8_t red = (pixel & red_mask) >> (8 + 8); // 这里的红色的掩码值是00ff0000，做与操作后，向右走16位，取出八位，也就是1byte的红色值

            *(RGB + t) = red;
            t++;

            uint8_t green = (pixel & green_mask) >> 8; // G绿色掩码值是0000ff00
            *(RGB + t) = green;
            t++;
            uint8_t blue = pixel & blue_mask; // 蓝色是000000ff
            *(RGB + t) = blue;
            t++;
        }
    }
    XDestroyImage(x_image);
    return RGB;
}
